#!/bin/bash

set -e

BOLD_RED='\033[1;31m'
GREEN='\033[32m'
BOLD_PURPLE='\033[1;35m'
NO_COLOR='\033[0m'

error() {
  echo -e "${BOLD_RED}>>> $1${NO_COLOR}"
}

info() {
  echo -e "${GREEN}>>> $1${NO_COLOR}"
}

error_exit() {
  error "$1"
  exit 1
}

PROVIDERS=("aws" "vb" "google")

usage() {
  echo -e "${BOLD_RED}>>> Usage: $0 [Machine Number] [Provider] [Options]${NO_COLOR}"
  echo -e "${BOLD_RED}>>> \tMachine Number:\t Must be at least 2${NO_COLOR}"
  echo -e "${BOLD_RED}>>> \tProvider:\t ${PROVIDERS[@]}${NO_COLOR}"
  echo -e "${BOLD_RED}>>> \tOptions:${NO_COLOR}"
  echo -e "${BOLD_RED}>>> \t -a, --alias-ip\t Add aws IP aliases(e.g., AlluxioMaster) into /etc/hosts${NO_COLOR}"
  exit
}

if [[ "$#" -lt 2 ]]; then
  usage
fi

TOTAL="$1"
[[ "${TOTAL}" -lt 2 ]] && error_exit "Machine number must be at least 2"
shift

ALLUXIO_MASTERS=$(grep Masters conf/alluxio.yml | cut -d':' -f 2 | tr -d ' ')
[[ "${TOTAL}" -lt "${ALLUXIO_MASTERS}" ]] && error_exit "Total number of machines must be at least the same with Masters in conf/alluxio.yml"

PROVIDER="$1"
[[ "${PROVIDERS[*]}" =~ (^| )$PROVIDER( |$) ]] || error_exit "Provider $PROVIDER is not supported yet"
shift

ALIAS_IP="false"
while [[ $# > 0 ]]; do
  key="$1"
  case $key in
    -a|--alias-ip)
      ALIAS_IP="true"
    ;;
    *)
    usage
    ;;
  esac
shift
done

HERE=`dirname "$0"`
cd $HERE

# make sure ec2.yml exists
if [[ "${PROVIDER}" == "aws" ]]; then
  if [[ ! -f conf/ec2.yml ]]; then
    error_exit "Please create your own conf/ec2.yml from conf/ec2.yml.template"
  fi
fi

# make sure gce.yml exists
if [[ "${PROVIDER}" == "google" ]]; then
  if [[ ! -f conf/gce.yml ]]; then
    error_exit "Please create your own conf/gce.yml from conf/gce.yml.template"
  fi
fi

cat > conf/init.yml <<EOF
# Generated by ../start, do not modify
Provider: $PROVIDER
MachineNumber: $TOTAL
EOF

mkdir -p ./files
# files/id_rsa, files/id_rsa.pub
if [[ ! -f ./files/id_rsa ]]; then
  ssh-keygen -f ./files/id_rsa -t rsa -N '' &>/dev/null
fi

# files/workers, order matters,
# when in Alluxio fault tolerant mode, the last one is often also chosen as Alluxio leader(not sure),
# when not in fault tolerant mode, the last one is Alluxio master,
# currently, for other systems like HDFS, fault tolerant mode hasn't been enabled, the last one is the single master.
rm -f files/workers
for i in `seq 1 $((TOTAL-ALLUXIO_MASTERS))`; do
  echo "AlluxioWorker${i}" >> files/workers
done
for ((i=2; i<=ALLUXIO_MASTERS; i++)); do
  echo "AlluxioMaster$i" >> files/workers
done
echo "AlluxioMaster" >> files/workers

# create virtual machine without provisioning
[[ "${PROVIDER}" == "aws" ]] && sgid=$(python bin/init_aws.py | tail -n 1)
# pass security group id to vagrant up
export AWS_SECURITY_GROUP_ID_TACH=$sgid
if [[ "${PROVIDER}" == "aws" && "$(python bin/is_aws_spot.py)" == "0" ]]; then
  # register handler to be run when exit 1
  trap "python bin/spot_request.py --cancel" ERR
  python bin/spot_request.py --submit --masters ${ALLUXIO_MASTERS}
  trap - ERR
else
  FULL_PROVIDER_NAME=$PROVIDER
  if [[ "${FULL_PROVIDER_NAME}" == "vb" ]]; then
    FULL_PROVIDER_NAME=virtualbox
    ./vbox/buildbox.sh
  fi
  VUP="vagrant up --provider=${FULL_PROVIDER_NAME} --no-provision"
  [[ ("${PROVIDER}" == "aws") || ("${PROVIDER}" == "google") ]] && VUP="${VUP} --parallel"
  $VUP
fi

# files/hosts
# IPs are automatically assigned by vagrant
# hack to find out the assigned IPs by looking into vagrant's ssh config to
# output of "ifconfig" is different in RHEL6.5 and RHEL7
#   inet:addr 127.0.0.1  VS. inet 127.0.0.1
# use "ip addr show" instead
rm -f files/hosts
for host in $(cat files/workers); do
  info "looking up IP of $host..."
  IP=$(vagrant ssh $host -c "ip addr show | grep -w inet | tail -1" | grep -w inet | awk '{print $2}' | cut -d'/' -f1)
  info "IP for $host is $IP"
  echo "$IP ${host}" >> files/hosts
done

# provision
vagrant provision

# IP aliasing
# For AWS and Google, the file ".vagrant/provisioners/ansible/inventory/vagrant_ansible_inventory"
# contains the hostname and IP information in the following format:
#
#   AlluxioWorker1 ansible_ssh_host=54.148.173.226 ansible_ssh_port=22
#   AlluxioMaster ansible_ssh_host=54.201.87.214 ansible_ssh_port=22
#
# The IP Aliasing is done by parsing this file to generate the "host ip" mapping

ANSIBLE_INFO_FILE=".vagrant/provisioners/ansible/inventory/vagrant_ansible_inventory"

if [[ "$ALIAS_IP" == "true" ]]; then
  TMP_FILE=./alluxio_hosts_tmp_file
  if [[ -f $ANSIBLE_INFO_FILE ]]; then
    echo -e "\n# Generated by Alluxio vagrant deploy script" > $TMP_FILE
    if [[ ( "${PROVIDER}" == "aws" ) || ( "${PROVIDER}" == "google" ) ]]; then
      regx="^Alluxio.*$"
      while read -r; do
        if [[ $REPLY =~ $regx ]]; then
          host=$(echo $REPLY | cut -d' ' -f1)
          ip=$(echo $REPLY | cut -d' ' -f2 | cut -d'=' -f2)
          echo "$host $ip" >>$TMP_FILE
        fi
      done < $ANSIBLE_INFO_FILE
    fi
    info "Adding following lines into /etc/hosts:"
    cat $TMP_FILE
    info ""
    sudo "$BASH" -c "sudo cat alluxio_hosts_tmp_file >> /etc/hosts"
    rm -f $TMP_FILE
  fi
fi

# show master IP
set +e # turn off the option that exits immediately when a command fails
if [[ ( "${PROVIDER}" == "aws" ) || ( "${PROVIDER}" == "google" ) ]]; then
  masters=$(cat $ANSIBLE_INFO_FILE | grep Master | cut -d' ' -f2 | cut -d'=' -f2)
  main_master=$(cat $ANSIBLE_INFO_FILE | grep 'AlluxioMaster ' | cut -d' ' -f2 | cut -d'=' -f2)
else
  masters=$(cat files/hosts | grep Master | cut -d' ' -f1)
  main_master=$(cat files/hosts | grep 'AlluxioMaster$' | cut -d' ' -f1)
fi
masters=($(echo $masters | tr '\n' ' ' | tr -d '\r'))

if [[ "$ALLUXIO_MASTERS" -gt 1 ]]; then
  info "Alluxio is in Fault Tolerant mode, waiting 20 seconds for leader to be selected and web UI to be available"
  sleep 20
  for master in ${masters[@]}; do
    info "Alluxio master: ${master}"
    info "Trying to curl Alluxio web UI at ${master}:19999"
    curl "${master}:19999" >/dev/null
    if [[ "$?" == "0" ]]; then
      info "Master public IP for Alluxio is ${master}, visit ${master}:19999 for Alluxio web UI"
      break
    fi
  done
  info "Master public IP for other softwares is ${main_master}"
else
  info "AlluxioMaster public IP is ${main_master}, visit ${main_master}:19999 for Alluxio web UI"
fi
info "visit default port of the web UI of what you deployed"
